"""检查会议数据集格式是否符合框架要求"""
import pandas as pd
import json
import sys
from pathlib import Path


def check_data_format(parquet_file):
    """
    检查数据格式是否满足框架需求
    
    必需字段：
    1. meeting_topic - Agent Loop使用（通过kwargs传递）
    2. required_attendees - Agent Loop使用
    3. optional_attendees - Agent Loop使用
    4. availability - Agent Loop使用
    5. historical_pattern - Agent Loop使用
    6. reward_model - NaiveRewardManager必需
       └─ ground_truth - NaiveRewardManager必需
    7. data_source - NaiveRewardManager使用（作为reward_fn_key）
    """
    
    print("="*80)
    print(f"检查数据文件: {parquet_file}")
    print("="*80)
    
    # 1. 检查文件是否存在
    if not Path(parquet_file).exists():
        print(f"❌ 错误: 文件不存在: {parquet_file}")
        return False
    
    # 2. 加载数据
    try:
        df = pd.read_parquet(parquet_file)
        print(f"✅ 成功加载数据文件")
        print(f"   数据量: {len(df)} 条")
    except Exception as e:
        print(f"❌ 错误: 无法加载文件: {e}")
        return False
    
    # 3. 检查必需字段
    required_fields = {
        'meeting_topic': '会议主题（Agent Loop使用）',
        'required_attendees': '必选人员列表（Agent Loop使用）',
        'optional_attendees': '可选人员列表（Agent Loop使用）',
        'availability': '人员日程可用性（Agent Loop使用）',
        'historical_pattern': '历史会议模式（Agent Loop使用）',
        'reward_model': '奖励模型字段（框架必需）',
        'data_source': '数据源标识（框架必需）',
    }
    
    print("\n" + "="*80)
    print("字段检查")
    print("="*80)
    
    all_fields_ok = True
    for field, desc in required_fields.items():
        if field in df.columns:
            print(f"✅ {field:25s} - {desc}")
        else:
            print(f"❌ {field:25s} - {desc} [缺失]")
            all_fields_ok = False
    
    if not all_fields_ok:
        print("\n❌ 数据格式检查失败：缺少必需字段")
        return False
    
    # 4. 检查 reward_model 结构
    print("\n" + "="*80)
    print("reward_model 结构检查")
    print("="*80)
    
    sample_reward_model = df['reward_model'].iloc[0]
    print(f"reward_model 类型: {type(sample_reward_model)}")
    
    if isinstance(sample_reward_model, dict):
        if 'ground_truth' in sample_reward_model:
            print(f"✅ reward_model.ground_truth 存在")
            print(f"   值: {repr(sample_reward_model['ground_truth'])}")
        else:
            print(f"❌ reward_model.ground_truth 缺失")
            print(f"   实际内容: {sample_reward_model}")
            return False
    else:
        # 可能是字符串（需要解析）
        try:
            if isinstance(sample_reward_model, str):
                parsed = json.loads(sample_reward_model)
                if 'ground_truth' in parsed:
                    print(f"✅ reward_model.ground_truth 存在（JSON字符串格式）")
                    print(f"   值: {repr(parsed['ground_truth'])}")
                else:
                    print(f"❌ reward_model.ground_truth 缺失")
                    return False
        except:
            print(f"❌ reward_model 格式错误，应为dict或JSON字符串")
            print(f"   实际类型: {type(sample_reward_model)}")
            print(f"   实际值: {sample_reward_model}")
            return False
    
    # 5. 检查数据样本
    print("\n" + "="*80)
    print("数据样本检查")
    print("="*80)
    
    sample = df.iloc[0]
    
    print(f"\n样本 0:")
    print(f"  meeting_topic: {sample['meeting_topic']}")
    print(f"  required_attendees: {sample['required_attendees']}")
    print(f"  optional_attendees: {sample['optional_attendees']}")
    print(f"  data_source: {repr(sample['data_source'])}")
    
    # 检查 availability 结构
    availability = sample['availability']
    if isinstance(availability, dict):
        print(f"  availability: dict with {len(availability)} people")
        first_person = list(availability.keys())[0]
        print(f"    示例: {first_person} -> {list(availability[first_person].keys())[:2]}...")
    else:
        print(f"  availability: {type(availability)}")
    
    # 检查 historical_pattern 结构
    historical = sample['historical_pattern']
    if isinstance(historical, dict):
        print(f"  historical_pattern: {historical.get('most_common_day')}, {historical.get('most_common_slot')}")
    else:
        print(f"  historical_pattern: {type(historical)}")
    
    # 6. 数据类型检查
    print("\n" + "="*80)
    print("数据类型检查")
    print("="*80)
    
    type_checks = [
        ('meeting_topic', str),
        ('required_attendees', (list, str)),  # 可能是list或JSON字符串
        ('optional_attendees', (list, str)),
        ('availability', (dict, str)),
        ('historical_pattern', (dict, str)),
        ('data_source', str),
    ]
    
    types_ok = True
    for field, expected_type in type_checks:
        actual_value = sample[field]
        actual_type = type(actual_value)
        
        if isinstance(expected_type, tuple):
            is_correct = isinstance(actual_value, expected_type)
            expected_str = ' or '.join([t.__name__ for t in expected_type])
        else:
            is_correct = isinstance(actual_value, expected_type)
            expected_str = expected_type.__name__
        
        if is_correct:
            print(f"✅ {field:25s}: {actual_type.__name__:15s} (期望: {expected_str})")
        else:
            print(f"⚠️  {field:25s}: {actual_type.__name__:15s} (期望: {expected_str})")
            # 不算错误，因为Parquet可能会序列化为字符串
    
    # 7. 统计信息
    print("\n" + "="*80)
    print("数据统计")
    print("="*80)
    
    print(f"总样本数: {len(df)}")
    print(f"总字段数: {len(df.columns)}")
    print(f"\n所有字段列表:")
    for i, col in enumerate(df.columns, 1):
        print(f"  {i:2d}. {col}")
    
    # 8. 与框架集成检查
    print("\n" + "="*80)
    print("框架集成检查")
    print("="*80)
    
    integration_checks = [
        ("data.prompt_key='meeting_topic'", 
         'meeting_topic' in df.columns,
         "配置中的prompt_key字段存在"),
        
        ("data.reward_fn_key='data_source'",
         'data_source' in df.columns,
         "配置中的reward_fn_key字段存在"),
        
        ("Agent Loop可获取required_attendees",
         'required_attendees' in df.columns,
         "Agent Loop通过kwargs['required_attendees']访问"),
        
        ("Agent Loop可获取availability",
         'availability' in df.columns,
         "Agent Loop通过kwargs['availability']访问"),
        
        ("NaiveRewardManager可获取reward_model",
         'reward_model' in df.columns,
         "通过data_item.non_tensor_batch['reward_model']访问"),
    ]
    
    for check_name, passed, description in integration_checks:
        status = "✅" if passed else "❌"
        print(f"{status} {check_name}")
        print(f"   说明: {description}")
    
    # 9. 最终结果
    print("\n" + "="*80)
    print("检查结果")
    print("="*80)
    
    if all_fields_ok:
        print("✅ 数据格式检查通过！")
        print("\n该数据集满足框架所有要求，可以用于训练。")
        return True
    else:
        print("❌ 数据格式检查失败！")
        print("\n请按照错误提示修复数据格式。")
        return False


def main():
    if len(sys.argv) < 2:
        print("用法: python check_data_format.py <parquet文件路径>")
        print("\n示例:")
        print("  python check_data_format.py train.parquet")
        print("  python check_data_format.py /path/to/meeting_dataset/train.parquet")
        sys.exit(1)
    
    parquet_file = sys.argv[1]
    success = check_data_format(parquet_file)
    
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()

